# 基于预训练StyleGAN2的树叶生成示例
from torchvision.utils import save_image

from skimage.measure import label
from PIL import Image
import numpy as np
import os
import matplotlib.pyplot as plt
import random
from skimage.transform import resize

def extract_leaf_instances(image_dir, mask_dir):
    """
    通过mask截取叶子实例并存储到list
    Args:
        image_dir: 叶子图像文件夹路径
        mask_dir: 对应mask文件夹路径
    Returns:
        包含所有叶子实例的list
    """
    leaf_instances = []

    # 获取mask文件列表
    mask_files = sorted(os.listdir(mask_dir))

    # 遍历mask文件
    for mask_file in mask_files:
        # 获取不带扩展名的文件名并删除_mask后缀
        base_name = os.path.splitext(mask_file)[0].replace('_mask', '')
        
        # 拼接对应的图像文件路径
        img_file = base_name + '.jpg'  # 假设图像文件扩展名为.jpg
        img_path = os.path.join(image_dir, img_file)
        
        # 如果图像文件不存在则跳过
        if not os.path.exists(img_path):
            print(img_path+'Image does not')
            continue
            
        # 加载图像和mask
        mask_path = os.path.join(mask_dir, mask_file)
        image = np.array(Image.open(img_path).convert('RGB'))
        mask = np.array(Image.open(mask_path).convert('L'))

        # 将mask二值化
        mask = (mask > 0).astype(np.uint8) * 1
        masked_image = image * mask[..., np.newaxis]
        
        # 调整叶片尺寸
        resized_instance = resize_leaf_instance(masked_image)
        leaf_instances.append(resized_instance)

    return leaf_instances


def save_leaf_instances(leaf_instances, save_dir):
    """保存叶子实例到指定目录
    Args:
        leaf_instances: 叶子实例列表
        save_dir: 保存目录路径
    """
    os.makedirs(save_dir, exist_ok=True)
    for idx, instance in enumerate(leaf_instances):
        save_path = os.path.join(save_dir, f'leaf_{idx}.png')
        Image.fromarray(instance).save(save_path)

def load_leaf_instances(load_dir):
    """从指定目录加载叶子实例
    Args:
        load_dir: 加载目录路径
    Returns:
        叶子实例列表
    """
    leaf_instances = []
    for file in sorted(os.listdir(load_dir)):
        if file.endswith('.png'):
            img_path = os.path.join(load_dir, file)
            instance = np.array(Image.open(img_path).convert('RGB'))
            leaf_instances.append(instance)
    return leaf_instances


def resize_leaf_instance(instance):
    """根据mask大小调整叶片尺寸并随机旋转
    Args:
        instance: 叶子实例(numpy数组)
    Returns:
        调整大小并旋转后的叶子实例
    """
    # 获取mask区域
    mask = instance[..., 0] > 0
    coords = np.where(mask)
    ymin, ymax = np.min(coords[0]), np.max(coords[0])
    xmin, xmax = np.min(coords[1]), np.max(coords[1])

    # 计算当前尺寸和mask像素点数量
    current_height = ymax - ymin
    current_width = xmax - xmin
    mask_pixel_count = np.sum(mask)  # 计算mask区域的像素点数量

    # 随机选择目标尺寸
    # random.seed(42)
    # target_size = random.randint(25, 35)
    target_size=25
    # 计算缩放比例（基于mask像素点数量）
    scale = (target_size * target_size) / mask_pixel_count

    # 裁剪并缩放
    cropped = instance[ymin:ymax, xmin:xmax]
    resized = resize(cropped, (int(current_height * scale), int(current_width * scale), 3),
                     preserve_range=True, anti_aliasing=True)
    
    # 随机旋转角度
    angle = random.uniform(-180, 180)
    
    # 将图像转换为PIL格式进行旋转
    resized_pil = Image.fromarray(resized.astype(np.uint8))
    rotated_pil = resized_pil.rotate(angle, expand=True, fillcolor=(0, 0, 0))
    
    # 转换回numpy数组
    rotated = np.array(rotated_pil)
    
    return rotated

def augment_apple_image(apple_image, apple_mask, leaf_instances, max_leaves=5):
    """增强苹果图像，随机添加叶子遮挡
    Args:
        apple_image: 苹果图像(numpy数组)
        apple_mask: 苹果mask(numpy数组)
        leaf_instances: 叶子实例列表
        max_leaves: 最大叶子数量
    Returns:
        增强后的图像和mask
    """
    # 随机选择叶子数量
    num_leaves = random.randint(1, max_leaves)
    
    # 随机选择叶子实例
    selected_leaves = random.sample(leaf_instances, num_leaves)
    
    # 复制原始图像和mask
    augmented_image = apple_image.copy()
    augmented_mask = apple_mask.copy()
    
    # 获取苹果mask区域
    apple_coords = np.where(apple_mask > 0)
    
    for leaf in selected_leaves:
        # 随机选择苹果上的位置
        idx = random.randint(0, len(apple_coords[0]) - 1)
        y, x = apple_coords[0][idx], apple_coords[1][idx]
        
        # 计算叶子放置位置
        h, w = leaf.shape[:2]
        y1, y2 = max(0, y - h//2), min(apple_image.shape[0], y + h//2 + h%2)
        x1, x2 = max(0, x - w//2), min(apple_image.shape[1], x + w//2 + w%2)
        
        # 调整叶子大小以适应边界
        leaf_h, leaf_w = y2 - y1, x2 - x1
        if leaf_h != h or leaf_w != w:
            leaf = resize(leaf, (leaf_h, leaf_w), preserve_range=True, anti_aliasing=True)
            leaf = leaf.astype(np.uint8)
        
        # 将叶子叠加到苹果图像上
        leaf_mask = leaf[..., 0] > 0
        augmented_image[y1:y2, x1:x2][leaf_mask] = leaf[leaf_mask]
        augmented_mask[y1:y2, x1:x2][leaf_mask] = 0  # 遮挡苹果mask
        
    return augmented_image, augmented_mask
# 在main代码中使用
if __name__ == '__main__':
    # 假设路径
    image_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\leafimg'
    mask_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\leafmask1'
    save_dir = r'D:\Jupyter\pytorch\d2l-zh\pytorch\appppppp\detection2\leaf_instances_sizereset'

    # 提取并保存叶子实例
    leaf_instances = extract_leaf_instances(image_dir, mask_dir)
    save_leaf_instances(leaf_instances, save_dir)

    # 加载叶子实例
    loaded_instances = load_leaf_instances(save_dir)
    print(loaded_instances[28].size)
    # 显示示例
    w=20
    h=20
    if len(loaded_instances) > 0:
        plt.figure(figsize=(w/100, h/100))
        plt.imshow(loaded_instances[28])
        plt.axis('off')
        plt.show()

